// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2016 Dmitry Vyukov <dvyukov@google.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_CXX11_THREADPOOL_NONBLOCKING_THREAD_POOL_H
#define EIGEN_CXX11_THREADPOOL_NONBLOCKING_THREAD_POOL_H

namespace Eigen {

template <typename Environment> class ThreadPoolTempl : public Eigen::ThreadPoolInterface
{
public:
    typedef typename Environment::Task Task;
    typedef RunQueue<Task, 1024> Queue;

    ThreadPoolTempl(int num_threads, Environment env = Environment()) : ThreadPoolTempl(num_threads, true, env) {}

    ThreadPoolTempl(int num_threads, bool allow_spinning, Environment env = Environment())
        : env_(env), num_threads_(num_threads), allow_spinning_(allow_spinning), thread_data_(num_threads), all_coprimes_(num_threads), waiters_(num_threads),
          global_steal_partition_(EncodePartition(0, num_threads_)), blocked_(0), spinning_(0), done_(false), cancelled_(false), ec_(waiters_)
    {
        waiters_.resize(num_threads_);
        // Calculate coprimes of all numbers [1, num_threads].
        // Coprimes are used for random walks over all threads in Steal
        // and NonEmptyQueueIndex. Iteration is based on the fact that if we take
        // a random starting thread index t and calculate num_threads - 1 subsequent
        // indices as (t + coprime) % num_threads, we will cover all threads without
        // repetitions (effectively getting a presudo-random permutation of thread
        // indices).
        eigen_plain_assert(num_threads_ < kMaxThreads);
        for (int i = 1; i <= num_threads_; ++i)
        {
            all_coprimes_.emplace_back(i);
            ComputeCoprimes(i, &all_coprimes_.back());
        }
#ifndef EIGEN_THREAD_LOCAL
        init_barrier_.reset(new Barrier(num_threads_));
#endif
        thread_data_.resize(num_threads_);
        for (int i = 0; i < num_threads_; i++)
        {
            SetStealPartition(i, EncodePartition(0, num_threads_));
            thread_data_[i].thread.reset(env_.CreateThread([this, i]() { WorkerLoop(i); }));
        }
#ifndef EIGEN_THREAD_LOCAL
        // Wait for workers to initialize per_thread_map_. Otherwise we might race
        // with them in Schedule or CurrentThreadId.
        init_barrier_->Wait();
#endif
    }

    ~ThreadPoolTempl()
    {
        done_ = true;

        // Now if all threads block without work, they will start exiting.
        // But note that threads can continue to work arbitrary long,
        // block, submit new work, unblock and otherwise live full life.
        if (!cancelled_)
        {
            ec_.Notify(true);
        }
        else
        {
            // Since we were cancelled, there might be entries in the queues.
            // Empty them to prevent their destructor from asserting.
            for (size_t i = 0; i < thread_data_.size(); i++) { thread_data_[i].queue.Flush(); }
        }
        // Join threads explicitly (by destroying) to avoid destruction order within
        // this class.
        for (size_t i = 0; i < thread_data_.size(); ++i) thread_data_[i].thread.reset();
    }

    void SetStealPartitions(const std::vector<std::pair<unsigned, unsigned>>& partitions)
    {
        eigen_plain_assert(partitions.size() == static_cast<std::size_t>(num_threads_));

        // Pass this information to each thread queue.
        for (int i = 0; i < num_threads_; i++)
        {
            const auto& pair = partitions[i];
            unsigned start = pair.first, end = pair.second;
            AssertBounds(start, end);
            unsigned val = EncodePartition(start, end);
            SetStealPartition(i, val);
        }
    }

    void Schedule(std::function<void()> fn) EIGEN_OVERRIDE { ScheduleWithHint(std::move(fn), 0, num_threads_); }

    void ScheduleWithHint(std::function<void()> fn, int start, int limit) override
    {
        Task t = env_.CreateTask(std::move(fn));
        PerThread* pt = GetPerThread();
        if (pt->pool == this)
        {
            // Worker thread of this pool, push onto the thread's queue.
            Queue& q = thread_data_[pt->thread_id].queue;
            t = q.PushFront(std::move(t));
        }
        else
        {
            // A free-standing thread (or worker of another pool), push onto a random
            // queue.
            eigen_plain_assert(start < limit);
            eigen_plain_assert(limit <= num_threads_);
            int num_queues = limit - start;
            int rnd = Rand(&pt->rand) % num_queues;
            eigen_plain_assert(start + rnd < limit);
            Queue& q = thread_data_[start + rnd].queue;
            t = q.PushBack(std::move(t));
        }
        // Note: below we touch this after making w available to worker threads.
        // Strictly speaking, this can lead to a racy-use-after-free. Consider that
        // Schedule is called from a thread that is neither main thread nor a worker
        // thread of this pool. Then, execution of w directly or indirectly
        // completes overall computations, which in turn leads to destruction of
        // this. We expect that such scenario is prevented by program, that is,
        // this is kept alive while any threads can potentially be in Schedule.
        if (!t.f)
        {
            ec_.Notify(false);
        }
        else
        {
            env_.ExecuteTask(t);  // Push failed, execute directly.
        }
    }

    void Cancel() EIGEN_OVERRIDE
    {
        cancelled_ = true;
        done_ = true;

        // Let each thread know it's been cancelled.
#ifdef EIGEN_THREAD_ENV_SUPPORTS_CANCELLATION
        for (size_t i = 0; i < thread_data_.size(); i++) { thread_data_[i].thread->OnCancel(); }
#endif

        // Wake up the threads without work to let them exit on their own.
        ec_.Notify(true);
    }

    int NumThreads() const EIGEN_FINAL { return num_threads_; }

    int CurrentThreadId() const EIGEN_FINAL
    {
        const PerThread* pt = const_cast<ThreadPoolTempl*>(this)->GetPerThread();
        if (pt->pool == this)
        {
            return pt->thread_id;
        }
        else
        {
            return -1;
        }
    }

private:
    // Create a single atomic<int> that encodes start and limit information for
    // each thread.
    // We expect num_threads_ < 65536, so we can store them in a single
    // std::atomic<unsigned>.
    // Exposed publicly as static functions so that external callers can reuse
    // this encode/decode logic for maintaining their own thread-safe copies of
    // scheduling and steal domain(s).
    static const int kMaxPartitionBits = 16;
    static const int kMaxThreads = 1 << kMaxPartitionBits;

    inline unsigned EncodePartition(unsigned start, unsigned limit) { return (start << kMaxPartitionBits) | limit; }

    inline void DecodePartition(unsigned val, unsigned* start, unsigned* limit)
    {
        *limit = val & (kMaxThreads - 1);
        val >>= kMaxPartitionBits;
        *start = val;
    }

    void AssertBounds(int start, int end)
    {
        eigen_plain_assert(start >= 0);
        eigen_plain_assert(start < end);  // non-zero sized partition
        eigen_plain_assert(end <= num_threads_);
    }

    inline void SetStealPartition(size_t i, unsigned val) { thread_data_[i].steal_partition.store(val, std::memory_order_relaxed); }

    inline unsigned GetStealPartition(int i) { return thread_data_[i].steal_partition.load(std::memory_order_relaxed); }

    void ComputeCoprimes(int N, MaxSizeVector<unsigned>* coprimes)
    {
        for (int i = 1; i <= N; i++)
        {
            unsigned a = i;
            unsigned b = N;
            // If GCD(a, b) == 1, then a and b are coprimes.
            while (b != 0)
            {
                unsigned tmp = a;
                a = b;
                b = tmp % b;
            }
            if (a == 1)
            {
                coprimes->push_back(i);
            }
        }
    }

    typedef typename Environment::EnvThread Thread;

    struct PerThread
    {
        constexpr PerThread() : pool(NULL), rand(0), thread_id(-1) {}
        ThreadPoolTempl* pool;  // Parent pool, or null for normal threads.
        uint64_t rand;          // Random generator state.
        int thread_id;          // Worker thread index in pool.
#ifndef EIGEN_THREAD_LOCAL
        // Prevent false sharing.
        char pad_[128];
#endif
    };

    struct ThreadData
    {
        constexpr ThreadData() : thread(), steal_partition(0), queue() {}
        std::unique_ptr<Thread> thread;
        std::atomic<unsigned> steal_partition;
        Queue queue;
    };

    Environment env_;
    const int num_threads_;
    const bool allow_spinning_;
    MaxSizeVector<ThreadData> thread_data_;
    MaxSizeVector<MaxSizeVector<unsigned>> all_coprimes_;
    MaxSizeVector<EventCount::Waiter> waiters_;
    unsigned global_steal_partition_;
    std::atomic<unsigned> blocked_;
    std::atomic<bool> spinning_;
    std::atomic<bool> done_;
    std::atomic<bool> cancelled_;
    EventCount ec_;
#ifndef EIGEN_THREAD_LOCAL
    std::unique_ptr<Barrier> init_barrier_;
    std::mutex per_thread_map_mutex_;  // Protects per_thread_map_.
    std::unordered_map<uint64_t, std::unique_ptr<PerThread>> per_thread_map_;
#endif

    // Main worker thread loop.
    void WorkerLoop(int thread_id)
    {
#ifndef EIGEN_THREAD_LOCAL
        std::unique_ptr<PerThread> new_pt(new PerThread());
        per_thread_map_mutex_.lock();
        bool insertOK = per_thread_map_.emplace(GlobalThreadIdHash(), std::move(new_pt)).second;
        eigen_plain_assert(insertOK);
        EIGEN_UNUSED_VARIABLE(insertOK);
        per_thread_map_mutex_.unlock();
        init_barrier_->Notify();
        init_barrier_->Wait();
#endif
        PerThread* pt = GetPerThread();
        pt->pool = this;
        pt->rand = GlobalThreadIdHash();
        pt->thread_id = thread_id;
        Queue& q = thread_data_[thread_id].queue;
        EventCount::Waiter* waiter = &waiters_[thread_id];
        // TODO(dvyukov,rmlarsen): The time spent in NonEmptyQueueIndex() is
        // proportional to num_threads_ and we assume that new work is scheduled at
        // a constant rate, so we set spin_count to 5000 / num_threads_. The
        // constant was picked based on a fair dice roll, tune it.
        const int spin_count = allow_spinning_ && num_threads_ > 0 ? 5000 / num_threads_ : 0;
        if (num_threads_ == 1)
        {
            // For num_threads_ == 1 there is no point in going through the expensive
            // steal loop. Moreover, since NonEmptyQueueIndex() calls PopBack() on the
            // victim queues it might reverse the order in which ops are executed
            // compared to the order in which they are scheduled, which tends to be
            // counter-productive for the types of I/O workloads the single thread
            // pools tend to be used for.
            while (!cancelled_)
            {
                Task t = q.PopFront();
                for (int i = 0; i < spin_count && !t.f; i++)
                {
                    if (!cancelled_.load(std::memory_order_relaxed))
                    {
                        t = q.PopFront();
                    }
                }
                if (!t.f)
                {
                    if (!WaitForWork(waiter, &t))
                    {
                        return;
                    }
                }
                if (t.f)
                {
                    env_.ExecuteTask(t);
                }
            }
        }
        else
        {
            while (!cancelled_)
            {
                Task t = q.PopFront();
                if (!t.f)
                {
                    t = LocalSteal();
                    if (!t.f)
                    {
                        t = GlobalSteal();
                        if (!t.f)
                        {
                            // Leave one thread spinning. This reduces latency.
                            if (allow_spinning_ && !spinning_ && !spinning_.exchange(true))
                            {
                                for (int i = 0; i < spin_count && !t.f; i++)
                                {
                                    if (!cancelled_.load(std::memory_order_relaxed))
                                    {
                                        t = GlobalSteal();
                                    }
                                    else
                                    {
                                        return;
                                    }
                                }
                                spinning_ = false;
                            }
                            if (!t.f)
                            {
                                if (!WaitForWork(waiter, &t))
                                {
                                    return;
                                }
                            }
                        }
                    }
                }
                if (t.f)
                {
                    env_.ExecuteTask(t);
                }
            }
        }
    }

    // Steal tries to steal work from other worker threads in the range [start,
    // limit) in best-effort manner.
    Task Steal(unsigned start, unsigned limit)
    {
        PerThread* pt = GetPerThread();
        const size_t size = limit - start;
        unsigned r = Rand(&pt->rand);
        // Reduce r into [0, size) range, this utilizes trick from
        // https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
        eigen_plain_assert(all_coprimes_[size - 1].size() < (1 << 30));
        unsigned victim = ((uint64_t)r * (uint64_t)size) >> 32;
        unsigned index = ((uint64_t)all_coprimes_[size - 1].size() * (uint64_t)r) >> 32;
        unsigned inc = all_coprimes_[size - 1][index];

        for (unsigned i = 0; i < size; i++)
        {
            eigen_plain_assert(start + victim < limit);
            Task t = thread_data_[start + victim].queue.PopBack();
            if (t.f)
            {
                return t;
            }
            victim += inc;
            if (victim >= size)
            {
                victim -= size;
            }
        }
        return Task();
    }

    // Steals work within threads belonging to the partition.
    Task LocalSteal()
    {
        PerThread* pt = GetPerThread();
        unsigned partition = GetStealPartition(pt->thread_id);
        // If thread steal partition is the same as global partition, there is no
        // need to go through the steal loop twice.
        if (global_steal_partition_ == partition)
            return Task();
        unsigned start, limit;
        DecodePartition(partition, &start, &limit);
        AssertBounds(start, limit);

        return Steal(start, limit);
    }

    // Steals work from any other thread in the pool.
    Task GlobalSteal() { return Steal(0, num_threads_); }

    // WaitForWork blocks until new work is available (returns true), or if it is
    // time to exit (returns false). Can optionally return a task to execute in t
    // (in such case t.f != nullptr on return).
    bool WaitForWork(EventCount::Waiter* waiter, Task* t)
    {
        eigen_plain_assert(!t->f);
        // We already did best-effort emptiness check in Steal, so prepare for
        // blocking.
        ec_.Prewait();
        // Now do a reliable emptiness check.
        int victim = NonEmptyQueueIndex();
        if (victim != -1)
        {
            ec_.CancelWait();
            if (cancelled_)
            {
                return false;
            }
            else
            {
                *t = thread_data_[victim].queue.PopBack();
                return true;
            }
        }
        // Number of blocked threads is used as termination condition.
        // If we are shutting down and all worker threads blocked without work,
        // that's we are done.
        blocked_++;
        // TODO is blocked_ required to be unsigned?
        if (done_ && blocked_ == static_cast<unsigned>(num_threads_))
        {
            ec_.CancelWait();
            // Almost done, but need to re-check queues.
            // Consider that all queues are empty and all worker threads are preempted
            // right after incrementing blocked_ above. Now a free-standing thread
            // submits work and calls destructor (which sets done_). If we don't
            // re-check queues, we will exit leaving the work unexecuted.
            if (NonEmptyQueueIndex() != -1)
            {
                // Note: we must not pop from queues before we decrement blocked_,
                // otherwise the following scenario is possible. Consider that instead
                // of checking for emptiness we popped the only element from queues.
                // Now other worker threads can start exiting, which is bad if the
                // work item submits other work. So we just check emptiness here,
                // which ensures that all worker threads exit at the same time.
                blocked_--;
                return true;
            }
            // Reached stable termination state.
            ec_.Notify(true);
            return false;
        }
        ec_.CommitWait(waiter);
        blocked_--;
        return true;
    }

    int NonEmptyQueueIndex()
    {
        PerThread* pt = GetPerThread();
        // We intentionally design NonEmptyQueueIndex to steal work from
        // anywhere in the queue so threads don't block in WaitForWork() forever
        // when all threads in their partition go to sleep. Steal is still local.
        const size_t size = thread_data_.size();
        unsigned r = Rand(&pt->rand);
        unsigned inc = all_coprimes_[size - 1][r % all_coprimes_[size - 1].size()];
        unsigned victim = r % size;
        for (unsigned i = 0; i < size; i++)
        {
            if (!thread_data_[victim].queue.Empty())
            {
                return victim;
            }
            victim += inc;
            if (victim >= size)
            {
                victim -= size;
            }
        }
        return -1;
    }

    static EIGEN_STRONG_INLINE uint64_t GlobalThreadIdHash() { return std::hash<std::thread::id>()(std::this_thread::get_id()); }

    EIGEN_STRONG_INLINE PerThread* GetPerThread()
    {
#ifndef EIGEN_THREAD_LOCAL
        static PerThread dummy;
        auto it = per_thread_map_.find(GlobalThreadIdHash());
        if (it == per_thread_map_.end())
        {
            return &dummy;
        }
        else
        {
            return it->second.get();
        }
#else
        EIGEN_THREAD_LOCAL PerThread per_thread_;
        PerThread* pt = &per_thread_;
        return pt;
#endif
    }

    static EIGEN_STRONG_INLINE unsigned Rand(uint64_t* state)
    {
        uint64_t current = *state;
        // Update the internal state
        *state = current * 6364136223846793005ULL + 0xda3e39cb94b95bdbULL;
        // Generate the random output (using the PCG-XSH-RS scheme)
        return static_cast<unsigned>((current ^ (current >> 22)) >> (22 + (current >> 61)));
    }
};

typedef ThreadPoolTempl<StlThreadEnvironment> ThreadPool;

}  // namespace Eigen

#endif  // EIGEN_CXX11_THREADPOOL_NONBLOCKING_THREAD_POOL_H
